using System;
using System.Runtime.CompilerServices;
using System.Runtime.ExceptionServices;


namespace Framework
{
    public class FTaskCompletionSource : IAwaiter
    {
        private AwaiterStatus m_Status;
        private ExceptionDispatchInfo m_Exception;
        private Action m_Continuation;

        AwaiterStatus IAwaiter.Status => this.m_Status;

        bool IAwaiter.IsCompleted => this.m_Status.IsCompleted();

        void IAwaiter.GetResult()
        {
            switch (this.m_Status)
            {
                case AwaiterStatus.Succeeded:
                    return;
                case AwaiterStatus.Faulted:
                {
                    this.m_Exception?.Throw();
                    this.m_Exception = null;
                    return;        
                }
                case AwaiterStatus.Canceled:
                {
                    this.m_Exception?.Throw();
                    this.m_Exception = null;
                    throw new OperationCanceledException();
                }
                default:
                    throw new NotSupportedException("Task does not allow call GetResult directly when task not completed. Please use 'await'.");
            }
        }

        void ICriticalNotifyCompletion.UnsafeOnCompleted(Action continuation)
        {
            this.m_Continuation = continuation;
            if (this.m_Status.IsCompleted())
            {
                this.TryInvokeContinuation();
            }
        }

        void INotifyCompletion.OnCompleted(Action continuation)
        {
            ((ICriticalNotifyCompletion) this).UnsafeOnCompleted(continuation);
        }

        public FTask Task => new FTask(this);

        public void SetResult()
        {
            if (this.TrySetResult())
            {
                return;
            }
            throw new InvalidOperationException("TaskT_TransitionToFinal_AlreadyCompleted");
        }

        public void SetException(Exception e)
        {
            if (this.TrySetException(e))
            {
                return;
            }
            throw new InvalidOperationException("TaskT_TransitionToFinal_AlreadyCompleted");
        }

        public void SetCanceled()
        {
            if (this.TrySetCanceled())
            {
                return;
            }
            throw new InvalidOperationException("TaskT_TransitionToFinal_AlreadyCompleted");
        }

        public void SetCanceled(OperationCanceledException e)
        {
            if (this.TrySetCanceled(e))
            {
                return;
            }
            throw new InvalidOperationException("TaskT_TransitionToFinal_AlreadyCompleted");
        }

        public bool TrySetResult()
        {
            if (this.m_Status.IsCompleted())
            {
                return false;
            }
            this.m_Status = AwaiterStatus.Succeeded;
            this.TryInvokeContinuation(); 
            return true;
        }

        public bool TrySetException(Exception e)
        {
            if (this.m_Status.IsCompleted())
            {
                return false;
            }
            this.m_Status = AwaiterStatus.Faulted;
            this.m_Exception = ExceptionDispatchInfo.Capture(e);
            this.TryInvokeContinuation();
            return true;
        }

        public bool TrySetCanceled()
        {
            if (this.m_Status.IsCompleted())
            {
                return false;
            }
            this.m_Status = AwaiterStatus.Canceled;
            this.TryInvokeContinuation();
            return true;
        }

        public bool TrySetCanceled(OperationCanceledException e)
        {
            if (this.m_Status.IsCompleted())
            {
                return false;
            }
            this.m_Status = AwaiterStatus.Canceled;
            this.m_Exception = ExceptionDispatchInfo.Capture(e);
            this.TryInvokeContinuation();
            return true;
        }

        private void TryInvokeContinuation()
        {
            this.m_Continuation?.Invoke();
            this.m_Continuation = null;
        }
    }


    public class FTaskCompletionSource<T> : IAwaiter<T>
    {
        private AwaiterStatus m_Status;
        private ExceptionDispatchInfo m_Exception;
        private Action m_Continuation;
        private T m_Result;

        AwaiterStatus IAwaiter.Status => this.m_Status;

        bool IAwaiter.IsCompleted => this.m_Status.IsCompleted();

        T IAwaiter<T>.GetResult()
        {
            switch (this.m_Status)
            {
                case AwaiterStatus.Succeeded:
                    return this.m_Result;
                case AwaiterStatus.Faulted:
                {
                    this.m_Exception?.Throw();
                    this.m_Exception = null;
                    return default;        
                }
                case AwaiterStatus.Canceled:
                {
                    this.m_Exception?.Throw();
                    this.m_Exception = null;
                    throw new OperationCanceledException();
                }
                default:
                    throw new NotSupportedException("Task does not allow call GetResult directly when task not completed. Please use 'await'.");
            }
        }

        void IAwaiter.GetResult()
        {
            ((IAwaiter<T>) this).GetResult();
        }

        void ICriticalNotifyCompletion.UnsafeOnCompleted(Action continuation)
        {
            this.m_Continuation = continuation;
            if (this.m_Status.IsCompleted())
            {
                this.TryInvokeContinuation();
            }
        }

        void INotifyCompletion.OnCompleted(Action continuation)
        {
            ((ICriticalNotifyCompletion) this).UnsafeOnCompleted(continuation);
        }

        public FTask<T> Task => new FTask<T>(this);

        public void SetResult(T result)
        {
            if (this.TrySetResult(result))
            {
                return;
            }
            throw new InvalidOperationException("TaskT_TransitionToFinal_AlreadyCompleted");
        }

        public void SetException(Exception e)
        {
            if (this.TrySetException(e))
            {
                return;
            }
            throw new InvalidOperationException("TaskT_TransitionToFinal_AlreadyCompleted");
        }

        public void SetCanceled()
        {
            if (this.TrySetCanceled())
            {
                return;
            }
            throw new InvalidOperationException("TaskT_TransitionToFinal_AlreadyCompleted");
        }

        public void SetCanceled(OperationCanceledException e)
        {
            if (this.TrySetCanceled(e))
            {
                return;
            }
            throw new InvalidOperationException("TaskT_TransitionToFinal_AlreadyCompleted");
        }

        public bool TrySetResult(T result)
        {
            if (this.m_Status.IsCompleted())
            {
                return false;
            }
            this.m_Status = AwaiterStatus.Succeeded;
            this.m_Result = result;
            this.TryInvokeContinuation(); 
            return true;
        }

        public bool TrySetException(Exception e)
        {
            if (this.m_Status.IsCompleted())
            {
                return false;
            }
            this.m_Status = AwaiterStatus.Faulted;
            this.m_Exception = ExceptionDispatchInfo.Capture(e);
            this.TryInvokeContinuation();
            return true;
        }

        public bool TrySetCanceled()
        {
            if (this.m_Status.IsCompleted())
            {
                return false;
            }
            this.m_Status = AwaiterStatus.Canceled;
            this.TryInvokeContinuation();
            return true;
        }

        public bool TrySetCanceled(OperationCanceledException e)
        {
            if (this.m_Status.IsCompleted())
            {
                return false;
            }
            this.m_Status = AwaiterStatus.Canceled;
            this.m_Exception = ExceptionDispatchInfo.Capture(e);
            this.TryInvokeContinuation();
            return true;
        }

        private void TryInvokeContinuation()
        {
            this.m_Continuation?.Invoke();
            this.m_Continuation = null;
        }
    }
}